import numpy as np
from .builder import ClusteringTreeBuilder

class Forest(object):

    def add_trees(self, trees):
        """ Add new trees to the forest (useful for distributed training)
        :trees: trees
        
        """
        self.trees.extend(trees)


    def datapoints_leaf_trees(self, dataset):
        """Return the matrix of datapoints (row) and the leave's id they  ended up in each tree (column) (size: datapoints X trees)

        :dataset: Original Data (non synthetic)
        :returns: nparray withsimilarity X.X

        """
        #dataset = np.atleast_2d(dataset)
        try:
            dt_leaf_trees = np.array([]).reshape((dataset.shape[0],0))
        except:
            dt_leaf_trees = np.array([]).reshape((1,0))
            dataset = dataset.reshape(1)
        for tree in self.trees:
            dt_leaf_trees = np.column_stack((dt_leaf_trees, tree.data_down_tree(dataset)[1])) 
        return dt_leaf_trees


    def generate_similarity_matrix(self, X_trees=None):
        """Generate the similarity matrix from the datapointsXtrees(leave's id in each tree)  matrix. It is an NXN matrix where n is the number of datapoints 

        :X_trees: N X T matrix
        :returns: similarity

        """
        #Create similarity after retraining
        if X_trees is not None:
            similarity_matrix = np.ones([X_trees.shape[0], X_trees.shape[0]], float)
            for idx, row in enumerate(X_trees):
                dt_row ,column = self.calculate_data_point_similarity(X_trees[idx], X_trees)
                similarity_matrix[idx] = dt_row
            return similarity_matrix
        else:
            #new prection with first data point for starting a new similarity matrix, so the distance of a dt point to itself is 1
            return  np.ones([1,1], float)

    
    def calculate_data_point_similarity(self, data_point_leaf_results, X_trees):
        """Calculates the similarity of a new datapoint to the existings ones, i.e., returns
        the dissimilarity row and column to be added to the similarity matrix

        :X_trees_dp: @todo
        :returns: @todo
        """
        #Create a vector (line of the matrix to be add) for the new DT_Agent
        data_point_similarity_row = np.zeros([1, X_trees.shape [0]])
        for idx, tree_result in enumerate(X_trees.T):
            id_column_elements_same_leaf = np.where(tree_result == data_point_leaf_results.T[idx])[0]
            for element_idx in id_column_elements_same_leaf:
                data_point_similarity_row[0][element_idx] += 1
        data_point_similarity_row = np.divide(data_point_similarity_row, X_trees.shape[1])
        #the matrix is symmetric so add the Tranpose and 1 similarity with itself
        data_point_similarity_column = np.append(data_point_similarity_row, 1).T
        return data_point_similarity_row, data_point_similarity_column

    def fit(self, X, Y, n_trees, **kwargs):
        #Create Trees
        self.trees = []
        for tree_n in xrange(n_trees):
            tree_leaf = self._tree_builder().fit(first=True, *self._bootstrap(X, Y), **kwargs)
            self.trees.append(tree_leaf)
        return self

    def _bootstrap(self, X, Y):
        size = int(X.shape[0]*0.67)
        sample = np.random.randint(X.shape[0], size=size)
        return X[sample], Y[sample]


class ClusteringForest(Forest):

    #Class to define tree builder and sserts
    def _tree_builder(self):
        return ClusteringTreeBuilder()

    def fit(self,
            X, Y,
            n_trees,
            **kwargs):

        #assert X.ndim == 2
        assert Y.ndim == 2

        #Calls super class(Forest) and the method fit
        return super(ClusteringForest, self).fit(X, Y, n_trees, **kwargs)
